import argparse
import sys
import os, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.utils import to_undirected
from torch_scatter import scatter
from torch_geometric.data import ShaDowKHopSampler

from logger import Logger
from dataset import *
from data_utils import normalize, gen_normalized_adjs, to_sparse_tensor, \
    load_fixed_splits, rand_splits, get_gpu_memory_map, count_parameters, reindex_env
from eval import evaluate_full, eval_acc, eval_rocauc, eval_f1
from parse import parser_add_main_args
from model import *
import time

# 每个数据集需要的信息：
# train_env_num 训练集环境数量
# env_num 环境数量
# test_ood_idx 测试集 OOD 索引
# train_idx 训练集索引
# valid_idx  验证集索引
# test_in_idx 测试集索引
# 下面的需要预处理创建额外的环境特征
# env 环境标签
# x 特征矩阵
# edge_index 边索引
# y 标签



# NOTE: for consistent data splits, see data_utils.rand_train_test_idx
# 定义了一个 fix_seed 函数，通过固定各个库的随机种子来确保实验的可重复性。
def fix_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

### Parse args ###
# 通过 argparse 解析命令行参数，并调用 parser_add_main_args 来添加主要的参数。
parser = argparse.ArgumentParser(description='General Training Pipeline')
parser_add_main_args(parser)
args = parser.parse_args()
print(args)

# 设置随机种子
fix_seed(args.seed)

# 根据参数选择使用CPU还是GPU进行训练
if args.cpu:
    device = torch.device("cpu")
else:
    device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")

### Load and preprocess data ###
# 根据不同的数据集名称，加载相应的数据集，并进行预处理。
# multi-graph datasets, divide graphs into train/valid/test
if args.dataset == 'twitch':
    dataset = load_twitch_dataset(args.data_dir, train_num=3)
elif args.dataset == 'elliptic':
    dataset = load_elliptic_dataset(args.data_dir, train_num=5)
# single-graph datasets, divide nodes into train/valid/test
elif args.dataset == 'arxiv':
    dataset = load_arxiv_dataset(args.data_dir, train_num=3)
# synthetic datasets, add spurious node features
elif args.dataset in ('cora', 'citeseer', 'pubmed'):
    dataset = load_synthetic_dataset(args.data_dir, args.dataset, train_num=3, combine=args.combine_result)
else:
    raise ValueError('Invalid dataname')

# len(dataset.y.shape) == 1:
# dataset.y 是一个张量，它包含了节点或图的标签（根据数据集类型不同）。
# dataset.y.shape 返回这个张量的形状，是一个元组。
# 如果 len(dataset.y.shape) == 1，表示这个张量是一维的，即它只有一个维度。

# dataset.y = dataset.y.unsqueeze(1):
# unsqueeze(1) 会在 dataset.y 的第一个维度（从0开始计数）插入一个新维度。
# 这样做的结果是将一维张量转换为二维张量，其形状从 (n,) 变为 (n, 1)，即每个标签单独成为一个包含一个元素的数组。
# 这种操作通常用于确保数据形状的一致性，特别是在需要与其他二维张量（如特征矩阵）进行操作时。
if len(dataset.y.shape) == 1:
    dataset.y = dataset.y.unsqueeze(1)

# dataset.y.max().item():
# dataset.y.max() 返回张量 dataset.y 中的最大值（这是一个张量）。
# item() 将包含单个值的张量转换为一个 Python 标量值。
# 例如，如果 dataset.y 的最大值是 4，那么 dataset.y.max().item() 的结果是 4。

# dataset.y.max().item() + 1:
# 由于标签通常是从 0 开始的整数，这里加 1 是为了得到标签的数量（类别数）。例如，如果标签的最大值是 4，那么标签的数量应为 5（包含 0 到 4 共五个类别）。

# dataset.y.shape[1]:
# dataset.y.shape[1] 返回 dataset.y 在第二个维度上的大小。
# 如果 dataset.y 是二维张量（形状为 (n, 1)），dataset.y.shape[1] 的结果是 1。

# c = max(dataset.y.max().item() + 1, dataset.y.shape[1]):
# 这行代码的目的是确定标签的类别数 c。
# 它取 dataset.y.max().item() + 1 和 dataset.y.shape[1] 这两个值中的较大值。
# 如果 dataset.y 已经是二维的，且其第二维度大小（dataset.y.shape[1]）大于 dataset.y.max().item() + 1，则采用 dataset.y.shape[1] 作为类别数。
c = max(dataset.y.max().item() + 1, dataset.y.shape[1])
# dataset.x 是特征矩阵，通常每一行表示一个节点或图的特征向量。
# dataset.x.shape 返回这个特征矩阵的形状，是一个元组。
# dataset.x.shape[1] 返回特征矩阵的第二个维度的大小，即特征的数量（每个节点或图的特征向量的长度）。
# d = dataset.x.shape[1] 将特征的数量赋值给变量 d。
d = dataset.x.shape[1]
# dataset.num_nodes 返回数据集中节点的数量。
n = dataset.num_nodes

# {dataset.edge_index.size(1)} 会被替换为数据集中所有边的数量。
# dataset.edge_index 是表示边的索引张量，通常是形状为 [2, num_edges] 的二维张量。
# dataset.edge_index.size(1) 返回边的数量，即第二个维度的大小（边的数量）。
print(f"dataset {args.dataset}: all nodes {dataset.num_nodes} | edges {dataset.edge_index.size(1)} | "
      + f"classes {c} | feats {d}")
print(f"train nodes {dataset.train_idx.shape[0]} | valid nodes {dataset.valid_idx.shape[0]} | "
      f"test in nodes {dataset.test_in_idx.shape[0]}")

# 遍历 dataset.test_ood_idx 列表，打印每个 OOD（Out-Of-Distribution）测试集的节点数量。
# dataset.test_ood_idx[i].shape[0] 获取第 i 个 OOD 测试集中的节点数量。
# 构建包含 OOD 测试集信息的字符串 m 并打印。
# 打印环境数量 dataset.env_num 和训练环境数量 dataset.train_env_num。
m = ""
for i in range(len(dataset.test_ood_idx)):
    m += f"test ood{i+1} nodes {dataset.test_ood_idx[i].shape[0]} "
print(m)
print(f'[INFO] env numbers: {dataset.env_num} train env numbers: {dataset.train_env_num}')

### Load method ###
# 判断数据集是否是多标签分类任务，proteins 和 ppi 数据集是多标签分类任务。
is_multilabel = args.dataset in ('proteins', 'ppi')

# 初始化 CaNet 模型。
# 参数：
# d：输入特征的维度。
# c：输出类别的数量。
# args：包含各种参数的对象。
# device：指定计算设备（如 GPU 或 CPU）。
# 将模型移到指定的设备上。
model = CaNet(d, c, args, device).to(device)

# 根据数据集类型选择损失函数：
# 对于 elliptic 和 twitch 数据集，使用 BCEWithLogitsLoss 损失函数，适用于二分类或多标签分类任务。
# 对于其他数据集，使用 CrossEntropyLoss 损失函数，适用于多分类任务。
if args.dataset in ('elliptic', 'twitch'):
    criterion = nn.BCEWithLogitsLoss(reduction='mean')
else:
    criterion = nn.CrossEntropyLoss(reduction='mean')

# 根据数据集类型选择评估函数：
# 对于 twitch 数据集，使用 eval_rocauc 函数计算 ROC AUC。
# 对于 elliptic 数据集，使用 eval_f1 函数计算 F1 分数。
# 对于其他数据集，使用 eval_acc 函数计算准确率。
if args.dataset in ('twitch'):
    eval_func = eval_rocauc
elif args.dataset in ('elliptic'):
    eval_func = eval_f1
else:
    eval_func = eval_acc

# 初始化日志记录器 Logger，用于记录训练和评估结果。
# 参数：
# args.runs：运行次数。
# args：包含各种参数的对象。
logger = Logger(args.runs, args)

# 设置模型为训练模式。
# 打印模型结构。
model.train()
print('MODEL:', model)

# 初始化训练准确率和验证准确率的列表，用于记录每次训练和验证的结果。
tr_acc, val_acc = [], []

### Training loop ###
# 运行多次训练，每次训练前重置模型参数。
for run in range(args.runs):
    model.reset_parameters()
    # 使用 Adam 优化器，并设置学习率和权重衰减。
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    # 初始化 best_val 变量，用于记录最佳验证准确率。
    best_val = float('-inf')

    # 将数据集的特征矩阵、标签、边索引和环境变量迁移到指定设备上。
    dataset.x, dataset.y, dataset.edge_index, dataset.env = \
        dataset.x.to(device), dataset.y.to(device), dataset.edge_index.to(device), dataset.env.to(device)

    # 遍历每个训练 epoch
    for epoch in range(args.epochs):
        # 设置模型为训练模式。
        model.train()
        # 清空优化器的梯度。
        optimizer.zero_grad()
        # 计算损失值。
        loss = model.loss_compute(dataset, criterion, args)
        # 反向传播计算梯度。
        loss.backward()
        # 更新模型参数。
        optimizer.step()
        # 评估模型，计算训练、验证和测试准确率。
        result = evaluate_full(model, dataset, eval_func)
        # 记录当前运行的结果。
        logger.add_result(run, result)

        # 将训练和验证准确率添加到对应的列表中。
        tr_acc.append(result[0])
        #todo 这个地方改了，原先是val_acc.append(result[2])
        val_acc.append(result[1])

        # 每隔 args.display_step 个 epoch 打印一次训练进度：
        # 打印当前 epoch 数、损失值、训练准确率、验证准确率和测试准确率。
        # 对于每个 OOD 测试集，打印其准确率。
        # 打印当前运行的统计结果。
        if epoch % args.display_step == 0:
            m = f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Train: {100 * result[0]:.2f}%, Valid: {100 * result[1]:.2f}%, Test In: {100 * result[2]:.2f}% '
            for i in range(len(result)-3):
                m += f'Test OOD{i+1}: {100 * result[i+3]:.2f}% '
            print(m)
    logger.print_statistics(run)

# 打印所有运行的统计结果。
# 如果设置了保存结果选项，将结果输出到指定文件。
logger.print_statistics()
if args.store_result:
    logger.output(args)